package com.ruoyi.learn.java.algorithm.searching;

/**
 * 红黑树 (Red-Black Tree) 的 Java 实现
 * 通过颜色和旋转/变色操作维持近似平衡。
 * @param <T> 可比较的泛型类型
 */
public class RedBlackTree<T extends Comparable<T>> {

    // 颜色枚举
    private enum Color {
        RED, BLACK
    }

    // 节点内部类
    private class Node {
        T data;
        Node left;
        Node right;
        Node parent; // 红黑树通常需要父指针
        Color color;

        public Node(T data) {
            this.data = data;
            this.left = NIL; // 初始化指向 NIL
            this.right = NIL;
            this.parent = null;
            this.color = Color.RED; // 新插入节点默认为红色
        }
    }

    // NIL 节点 (代表空叶子)，是黑色的
    private final Node NIL = new Node(null);
    {
        NIL.color = Color.BLACK;
        NIL.left = null;
        NIL.right = null;
        NIL.parent = null;
    }

    // 树的根节点
    private Node root;

    // ------------------- 构造函数 -------------------
    public RedBlackTree() {
        this.root = NIL; // 初始为空树
    }

    // ------------------- 辅助方法 -------------------

    private void leftRotate(Node x) {
        Node y = x.right;
        x.right = y.left;
        if (y.left != NIL) {
            y.left.parent = x;
        }
        y.parent = x.parent;
        if (x.parent == null) {
            root = y;
        } else if (x == x.parent.left) {
            x.parent.left = y;
        } else {
            x.parent.right = y;
        }
        y.left = x;
        x.parent = y;
    }

    private void rightRotate(Node x) {
        Node y = x.left;
        x.left = y.right;
        if (y.right != NIL) {
            y.right.parent = x;
        }
        y.parent = x.parent;
        if (x.parent == null) {
            root = y;
        } else if (x == x.parent.right) {
            x.parent.right = y;
        } else {
            x.parent.left = y;
        }
        y.right = x;
        x.parent = y;
    }

    // 设置节点的父节点和左右孩子，并更新孩子节点的父指针
    private void transplant(Node u, Node v) {
        if (u.parent == null) {
            root = v;
        } else if (u == u.parent.left) {
            u.parent.left = v;
        } else {
            u.parent.right = v;
        }
        v.parent = u.parent;
    }

    // ------------------- 核心操作 -------------------

    public void insert(T data) {
        if (data == null) {
            throw new IllegalArgumentException("Cannot insert null data");
        }

        Node newNode = new Node(data);
        Node parent = null;
        Node current = root;

        // 1. 找到插入位置 (标准 BST 插入)
        while (current != NIL) {
            parent = current;
            int compareResult = data.compareTo(current.data);
            if (compareResult < 0) {
                current = current.left;
            } else if (compareResult > 0) {
                current = current.right;
            } else {
                // 相等，通常不允许重复，直接返回 (可按需修改)
                return;
            }
        }

        newNode.parent = parent;
        if (parent == null) {
            // 树为空，新节点成为根
            root = newNode;
        } else if (data.compareTo(parent.data) < 0) {
            parent.left = newNode;
        } else {
            parent.right = newNode;
        }

        // 新节点的左右孩子指向 NIL
        newNode.left = NIL;
        newNode.right = NIL;

        // 2. 修复红黑树性质 (插入修复)
        insertFixup(newNode);
    }

    // 插入修复
    private void insertFixup(Node z) {
        // 只要 z 不是根节点且其父节点是红色，就需要修复
        while (z != root && z.parent.color == Color.RED) {
            if (z.parent == z.parent.parent.left) {
                // 父节点是祖父节点的左孩子
                Node y = z.parent.parent.right; // 叔叔节点
                if (y.color == Color.RED) {
                    // Case 1: 叔叔是红色
                    z.parent.color = Color.BLACK;
                    y.color = Color.BLACK;
                    z.parent.parent.color = Color.RED;
                    z = z.parent.parent; // 将问题上移
                } else {
                    // 叔叔是黑色
                    if (z == z.parent.right) {
                        // Case 2: z 是右孩子 (需要先左旋转)
                        z = z.parent;
                        leftRotate(z);
                    }
                    // Case 3: z 是左孩子
                    z.parent.color = Color.BLACK;
                    z.parent.parent.color = Color.RED;
                    rightRotate(z.parent.parent);
                }
            } else {
                // 对称情况 (父节点是右孩子)
                Node y = z.parent.parent.left; // 叔叔节点
                if (y.color == Color.RED) {
                    // Case 1: 叔叔是红色
                    z.parent.color = Color.BLACK;
                    y.color = Color.BLACK;
                    z.parent.parent.color = Color.RED;
                    z = z.parent.parent;
                } else {
                    if (z == z.parent.left) {
                        // Case 2: z 是左孩子 (需要先右旋转)
                        z = z.parent;
                        rightRotate(z);
                    }
                    // Case 3: z 是右孩子
                    z.parent.color = Color.BLACK;
                    z.parent.parent.color = Color.RED;
                    leftRotate(z.parent.parent);
                }
            }
        }
        // 确保根节点是黑色
        root.color = Color.BLACK;
    }

    public void delete(T data) {
        if (data == null) {
            throw new IllegalArgumentException("Cannot delete null data");
        }

        Node z = searchNode(root, data);
        if (z == null || z == NIL) {
            return; // 节点不存在
        }

        Node x;
        Node y = z; // y 是真正被删除或移动的节点
        Color yOriginalColor = y.color;

        if (z.left == NIL) {
            // 情况 1: z 没有左孩子
            x = z.right;
            transplant(z, z.right);
        } else if (z.right == NIL) {
            // 情况 2: z 没有右孩子
            x = z.left;
            transplant(z, z.left);
        } else {
            // 情况 3: z 有两个孩子
            // 找到 z 的后继 (右子树的最小值)
            y = findMin(z.right);
            yOriginalColor = y.color;
            x = y.right; // y 的后继要么是 NIL，要么是右孩子 (因为 y 没有左孩子)
            if (y.parent == z) {
                x.parent = y; // 特殊情况
            } else {
                transplant(y, y.right);
                y.right = z.right;
                y.right.parent = y;
            }
            transplant(z, y);
            y.left = z.left;
            y.left.parent = y;
            y.color = z.color; // 继承原节点的颜色
        }

        // 如果被删除节点的原始颜色是黑色，需要修复
        if (yOriginalColor == Color.BLACK) {
            deleteFixup(x);
        }
    }

    // 删除修复
    private void deleteFixup(Node x) {
        // x 是一个额外的"黑色"，需要通过旋转和变色将其"推"到根节点或消除
        while (x != root && x.color == Color.BLACK) {
            if (x == x.parent.left) {
                Node w = x.parent.right; // 兄弟节点
                if (w.color == Color.RED) {
                    // Case 1: 兄弟是红色
                    w.color = Color.BLACK;
                    x.parent.color = Color.RED;
                    leftRotate(x.parent);
                    w = x.parent.right; // 更新兄弟
                }
                if (w.left.color == Color.BLACK && w.right.color == Color.BLACK) {
                    // Case 2: 兄弟是黑色，且两个侄子都是黑色
                    w.color = Color.RED;
                    x = x.parent; // 将额外的黑色上移
                } else {
                    if (w.right.color == Color.BLACK) {
                        // Case 3: 兄弟是黑色，左侄红，右侄黑
                        w.left.color = Color.BLACK;
                        w.color = Color.RED;
                        rightRotate(w);
                        w = x.parent.right; // 更新兄弟
                    }
                    // Case 4: 兄弟是黑色，右侄红
                    w.color = x.parent.color;
                    x.parent.color = Color.BLACK;
                    w.right.color = Color.BLACK;
                    leftRotate(x.parent);
                    x = root; // 修复完成
                }
            } else {
                // 对称情况 (x 是右孩子)
                Node w = x.parent.left; // 兄弟节点
                if (w.color == Color.RED) {
                    w.color = Color.BLACK;
                    x.parent.color = Color.RED;
                    rightRotate(x.parent);
                    w = x.parent.left;
                }
                if (w.right.color == Color.BLACK && w.left.color == Color.BLACK) {
                    w.color = Color.RED;
                    x = x.parent;
                } else {
                    if (w.left.color == Color.BLACK) {
                        w.right.color = Color.BLACK;
                        w.color = Color.RED;
                        leftRotate(w);
                        w = x.parent.left;
                    }
                    w.color = x.parent.color;
                    x.parent.color = Color.BLACK;
                    w.left.color = Color.BLACK;
                    rightRotate(x.parent);
                    x = root;
                }
            }
        }
        x.color = Color.BLACK; // 确保根节点是黑色
    }

    // 在树中搜索包含 data 的节点
    private Node searchNode(Node node, T data) {
        while (node != NIL && data.compareTo(node.data) != 0) {
            if (data.compareTo(node.data) < 0) {
                node = node.left;
            } else {
                node = node.right;
            }
        }
        return node;
    }

    public boolean search(T data) {
        if (data == null) {
            return false;
        }
        return searchNode(root, data) != NIL;
    }

    private Node findMin(Node node) {
        while (node.left != NIL) {
            node = node.left;
        }
        return node;
    }

    public T findMin() {
        if (root == NIL) {
            return null;
        }
        return findMin(root).data;
    }

    private Node findMax(Node node) {
        while (node.right != NIL) {
            node = node.right;
        }
        return node;
    }

    public T findMax() {
        if (root == NIL) {
            return null;
        }
        return findMax(root).data;
    }

    // ------------------- 遍历操作 -------------------

    public void inOrderTraversal() {
        System.out.print("In-order: ");
        inOrderRecursive(root);
        System.out.println();
    }

    private void inOrderRecursive(Node node) {
        if (node != NIL) {
            inOrderRecursive(node.left);
            System.out.print(node.data + "(" + (node.color == Color.RED ? "R" : "B") + ") ");
            inOrderRecursive(node.right);
        }
    }

    public void preOrderTraversal() {
        System.out.print("Pre-order: ");
        preOrderRecursive(root);
        System.out.println();
    }

    private void preOrderRecursive(Node node) {
        if (node != NIL) {
            System.out.print(node.data + "(" + (node.color == Color.RED ? "R" : "B") + ") ");
            preOrderRecursive(node.left);
            preOrderRecursive(node.right);
        }
    }

    // ------------------- 辅助方法 -------------------

    public boolean isEmpty() {
        return root == NIL;
    }

    // 计算黑高 (从根到叶子的黑色节点数)
    public int blackHeight() {
        int bh = 0;
        Node node = root;
        while (node != NIL) {
            if (node.color == Color.BLACK) {
                bh++;
            }
            node = node.left; // 任意路径都可以
        }
        return bh;
    }

    public void clear() {
        root = NIL;
    }

    // ------------------- 主函数用于测试 -------------------
    public static void main(String[] args) {
        RedBlackTree<Integer> rbt = new RedBlackTree<>();

        int[] values = {10, 20, 30, 15, 25, 5, 1};
        System.out.println("Inserting values: ");
        for (int val : values) {
            System.out.print(val + " ");
            rbt.insert(val);
        }
        System.out.println();

        rbt.inOrderTraversal();  // 显示数据和颜色
        rbt.preOrderTraversal();

        System.out.println("Search 15: " + rbt.search(15)); // true
        System.out.println("Search 100: " + rbt.search(100)); // false

        System.out.println("Min: " + rbt.findMin()); // 1
        System.out.println("Max: " + rbt.findMax()); // 30
        System.out.println("Black Height: " + rbt.blackHeight());

        System.out.println("Deleting 20");
        rbt.delete(20);
        rbt.inOrderTraversal();
        rbt.preOrderTraversal();
        System.out.println("Black Height after deletion: " + rbt.blackHeight());
    }
}
